import numpy as np
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeRegressor

# 产生一个随机数
rng=np.random.RandomState(1)
X=np.sort(10*rng.rand(80,1),axis=0)
Y=np.sin(X).ravel()
Y[::5]+=3*(0.5-rng.rand(16))
#构建不同深度的决策树
clf_0=DecisionTreeRegressor(max_depth=1)
clf_1=DecisionTreeRegressor(max_depth=2)
clf_2=DecisionTreeRegressor(max_depth=3)
clf_3=DecisionTreeRegressor(max_depth=5)
clf_0.fit(X,Y)
clf_1.fit(X,Y)
clf_2.fit(X,Y)
clf_3.fit(X,Y)
#创建预测模拟数据
X_test=np.arange(0.0,10,0.01)[:,np.newaxis]
Y_0=clf_0.predict(X_test)
Y_1=clf_1.predict(X_test)
Y_2=clf_2.predict(X_test)
Y_3=clf_3.predict(X_test)
#图表展示
plt.figure(figsize=(16,9),dpi=80,facecolor='w')
plt.scatter(X,Y,c='k',s=10,label='data')
plt.plot(X_test,Y_0,c='y',lw=2,label='max_depth=1,$R^2$=%.3f'%(clf_0.score(X,Y)))
plt.plot(X_test,Y_1,c='g',lw=2,label='max_depth=2,$R^2$=%.3f'%(clf_1.score(X,Y)))
plt.plot(X_test,Y_2,c='r',lw=2,label='max_depth=3,$R^2$=%.3f'%(clf_2.score(X,Y)))
plt.plot(X_test,Y_3,c='b',lw=2,label='max_depth=5,$R^2$=%.3f'%(clf_3.score(X,Y)))
plt.xlabel('X',horizontalalignment='left')
plt.ylabel('Y')
plt.title('Decision Tree Regression')
plt.legend()
plt.show()